/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: ai stats log manager class
 */

#include "ai_stats_log_file_writer.h"

#include <algorithm> // sort
#include <thread>
#include <algorithm>
#include <string>

#include "infra/base/scope_guard.h"
#include "infra/base/assertion.h"
#include "infra/om/stats/ai_stats_log_utils.h"
#ifdef STATS_CLIENT
#include "infra/memory_manager/client/fd_manager_ddk.h"
#include "infra/rpc/rpc_request_client.h"
#endif

namespace hiai {

namespace {
    bool g_isDestroyed = false;
    std::mutex g_resourceMutex;
}
/* **********************************Start AiStatsLogFileWriter************************* */
AiStatsLogFileWriter* AiStatsLogFileWriter::GetInstance()
{
    static AiStatsLogFileWriter instance;
    return &instance;
}

void AiStatsLogFileWriter::DoAsyncWrite()
{
    std::lock_guard<std::mutex> lk(g_resourceMutex);
    // 如果已经进入析构函数，则不再执行任何逻辑（创建线程）
    if (!g_isDestroyed) {
        threadCnt_++;
        // 创建线程前加计数，保证在析构中进入yield等待
        pthread_t writeFileThread;
        int ret = pthread_create(&writeFileThread, nullptr, &AiStatsLogFileWriter::Start, this);
        if (ret != 0) {
            threadCnt_--;
            return;
        }

        pthread_detach(writeFileThread);
    }
}

/*lint -e1744*/
AiStatsLogFileWriter::AiStatsLogFileWriter()
{
#ifdef STATS_CLIENT
    rpc::RpcRequestClient::GetInstance();
    FdManager::Init();
#endif
}

AiStatsLogFileWriter::~AiStatsLogFileWriter()
{
    g_isDestroyed = true;
    std::unique_lock<std::mutex> lk(g_resourceMutex);
    while (threadCnt_ != 0) {
        g_resourceMutex.unlock();
        std::this_thread::yield();
        g_resourceMutex.lock();
    }
}

void AiStatsLogFileWriter::CopyBuffer(
    std::map<pid_t, std::map<std::string, AiStatsData>> buffer, AiStatsEngineType type)
{
    hiai::ScopeGuard guard([this]() { DoAsyncWrite(); });
    {
        std::lock_guard<std::mutex> lk(this->engineBufferMapLock_);
        // 如果数据不存在，则添加
        auto iter = engineBufferMap_.find(type);
        if (iter == engineBufferMap_.end()) {
            engineBufferMap_[type] = buffer;
            return;
        }

        // 如果buffer中已经存在该engine的数据，则追加。
        for (auto it = buffer.begin(); it != buffer.end(); it++) {
            if (engineBufferMap_[type].find(it->first) == engineBufferMap_[type].end()) {
                engineBufferMap_[type][it->first] = it->second;
                continue;
            }
            for (auto iterm = it->second.begin(); iterm != it->second.end(); iterm++) {
                if (engineBufferMap_[type][it->first].find(iterm->first) == engineBufferMap_[type][it->first].end()) {
                    engineBufferMap_[type][it->first][iterm->first] = iterm->second;
                    continue;
                }
                engineBufferMap_[type][it->first][iterm->first] =
                    AiStatsLogUtils::AddAiStatsData(engineBufferMap_[type][it->first][iterm->first], iterm->second);
            }
        }
    }
    return;
}

#ifdef STATS_SERVER
void AiStatsLogFileWriter::WriteMessageToFile()
{
    for (std::pair<AiStatsEngineType, std::string> message : messageBuffer_) {
        FileWriter::Write(message.first, message.second);
    }
    messageBuffer_.clear();
}
void AiStatsLogFileWriter::WriteMessage(const AiStatsEngineType type, std::string message)
{
    {
        std::lock_guard<std::mutex> lk(this->engineBufferMapLock_);
        messageBuffer_.push_back(std::make_pair(type, message));
    }
    DoAsyncWrite();
}
#endif

void AiStatsLogFileWriter::WriteToFile(
    const AiStatsEngineType type, std::map<pid_t, std::map<std::string, AiStatsData>>& logs)
{
    std::string message;
    for (auto it = logs.cbegin(); it != logs.cend(); it++) {
        for (auto iter = it->second.cbegin(); iter != it->second.cend(); iter++) {
            message += AiStatsLogUtils::BuildString(iter->second);
            message += "\n";
        }
    }

    FileWriter::Write(type, message);
}

void* AiStatsLogFileWriter::Start(void* arg)
{
    auto* pWriter = static_cast<AiStatsLogFileWriter*>(arg);
    HIAI_EXPECT_NOT_NULL_R(pWriter, nullptr);

    hiai::ScopeGuard guard([pWriter]() {
        std::lock_guard<std::mutex> lk(g_resourceMutex);
        pWriter->threadCnt_--;
    });

    std::vector<AiStatsEngineType> engineTypeVect;

    std::unique_lock<std::mutex> threadLock(pWriter->engineBufferMapLock_);

    // 2. buffer不为空时，需要将其内容写入文件
    for (auto iter = pWriter->engineBufferMap_.begin(); iter != pWriter->engineBufferMap_.end(); ++iter) {
        AiStatsEngineType type = iter->first;
        pWriter->WriteToFile(type, iter->second);
        engineTypeVect.push_back(type);
    }
    // 3. 将已经写入文件的buffer数据删除
    for (const auto& mType : engineTypeVect) {
        pWriter->engineBufferMap_.erase(mType);
    }
    engineTypeVect.clear();

#ifdef STATS_SERVER
    pWriter->WriteMessageToFile();
#endif
    return nullptr;
}
/* **********************************end AiStatsLogFileWriter************************* */
} // namespace hiai
